import torch
from torch import nn
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader
from torch.functional import F

# 加载数据
transform = transforms.Compose([transforms.ToTensor()])
# 训练数据
trainset = datasets.MNIST('../data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=16, shuffle=True)

# 测试数据
testset = datasets.MNIST(root="../data", train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=16, shuffle=True)


# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.pooling = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(320, 10)

    def forward(self, x):
        batch_size = x.size(0)
        x = F.relu(self.conv1(x))  # batch * 1 * 28 * 28 ====> batch * 10 * 24 * 24
        x = self.pooling(x)  # batch * 10 * 24 * 24 ====> batch * 10 * 12 * 12
        x = F.relu(self.conv2(x))  #  batch * 10 * 12 * 12 ====> batch * 20 * 8 * 8
        x = self.pooling(x)  # batch * 1 * 8 * 8 ====> batch * 20 * 4 * 4
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x


# 创建模型
model = Net()

# 创建交叉熵损失
criterion = nn.CrossEntropyLoss(reduction="mean")

# 创建优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练
epochs = 10
def train(epoch):
    running_loss = 0.0
    for batch_index, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if batch_index % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_index + 1, running_loss / 300))
            running_loss = 0.0

# 测试
def test():
    total = 0
    correct = 0
    with torch.no_grad():
        for data in testloader:
            inputs, labels = data
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print((100 * correct / total), correct, total)

if __name__ == '__main__':
    for epoch in range(epochs):
        train(epoch)
        test()
